package org.future.task.redis.cluster;

import io.lettuce.core.RedisFuture;
import io.lettuce.core.RedisURI;
import io.lettuce.core.cluster.RedisClusterClient;
import io.lettuce.core.cluster.api.StatefulRedisClusterConnection;
import io.lettuce.core.cluster.api.async.RedisAdvancedClusterAsyncCommands;
import redis.clients.jedis.HostAndPort;

import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executor;
import java.util.concurrent.Semaphore;
import java.util.stream.Collectors;

/**
 * @author Dongwei-Chen
 * @Date 2023/12/13 10:10
 * @Description
 */
public class AsyncClusterManager<K, V> implements AsyncClusterCommand<K, V> {

    private static StatefulRedisClusterConnection<String, String> connect;

    public AsyncClusterManager(List<HostAndPort> nodes, String password) {
        List<RedisURI> redisURIS = nodes.stream().map(node -> {
            String host = node.getHost();
            int port = node.getPort();
            RedisURI.Builder builder = RedisURI.builder().withHost(host)
                    .withPort(port)
                    .withDatabase(0);
            if (password != null) {
                builder = builder.withPassword(password);
            }
            return builder.build();
        }).collect(Collectors.toList());
        RedisClusterClient client = RedisClusterClient.create(redisURIS);
        connect = client.connect();
        connect.setAutoFlushCommands(false);
    }

    public AsyncClusterManager(List<HostAndPort> nodes) {
        this(nodes, null);
    }

    @Override
    public void batchExecutorSync(List<K> batch, AsyncClusterAbstract<K, V> asyncCommandAbstract) {
        RedisAdvancedClusterAsyncCommands<String, String> async = connect.async();
        CountDownLatch downLatch = new CountDownLatch(batch.size());
        for (K row : batch) {
            try {
                RedisFuture<V> command = asyncCommandAbstract.command(async, row);
                command.whenComplete((v, throwable) -> {
                    try {
                        asyncCommandAbstract.future(row, v, throwable);
                    } catch (Exception e) {
                        e.printStackTrace();
                    } finally {
                        downLatch.countDown();
                    }
                });
            } catch (Throwable e) {
                downLatch.countDown();
            }
        }
        async.flushCommands();
        try {
            downLatch.await();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
    }

    @Override
    public void batchExecutorAsync(List<K> batch, AsyncClusterAbstract<K, V> asyncCommandAbstract, Executor executor) {
        RedisAdvancedClusterAsyncCommands<String, String> async = connect.async();
        CountDownLatch downLatch = new CountDownLatch(batch.size());
        for (K row : batch) {
            try {
                RedisFuture<V> command = asyncCommandAbstract.command(async, row);
                command.whenCompleteAsync((v, throwable) -> {
                    try {
                        asyncCommandAbstract.future(row, v, throwable);
                    } catch (Exception e) {
                        e.printStackTrace();
                    } finally {
                        downLatch.countDown();
                    }
                }, executor);
            } catch (Throwable e) {
                downLatch.countDown();
            }
        }
        async.flushCommands();
        try {
            downLatch.await();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
    }

    @Override
    public void batchExecutorAsync(List<K> batch, AsyncClusterAbstract<K, V> asyncCommandAbstract) {
        RedisAdvancedClusterAsyncCommands<String, String> async = connect.async();
        CountDownLatch downLatch = new CountDownLatch(batch.size());
        for (K row : batch) {
            try {
                RedisFuture<V> command = asyncCommandAbstract.command(async, row);
                command.whenCompleteAsync((v, throwable) -> {
                    try {
                        asyncCommandAbstract.future(row, v, throwable);
                    } catch (Exception e) {
                        e.printStackTrace();
                    } finally {
                        downLatch.countDown();
                    }
                });
            } catch (Throwable e) {
                downLatch.countDown();
            }
        }
        async.flushCommands();
        try {
            downLatch.await();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
    }


    @Override
    public void batchFlowSync(List<K> batch, AsyncClusterAbstract<K, V> asyncCommandAbstract) {
        RedisAdvancedClusterAsyncCommands<String, String> async = connect.async();
        Semaphore semaphore = new Semaphore(1000);
        for (K row : batch) {
            try {
                semaphore.acquire();
                RedisFuture<V> command = asyncCommandAbstract.command(async, row);
                command.whenComplete((v, throwable) -> {
                    try {
                        asyncCommandAbstract.future(row, v, throwable);
                    } catch (Exception e) {
                        e.printStackTrace();
                    } finally {
                        semaphore.release();
                    }
                });
                async.flushCommands();
            } catch (Throwable e) {
                semaphore.release();
            }
        }
    }
}
